import {
  type HostVulnerabilityCounts,
  type HostVulnerabilityMap,
} from "~/types/vulnerabilityTypes";
import { type VulnerabilityAggregationResponse } from "~/types/vulnerabilityTypes";

/**
 * Gets the vulnerability counts for a specific host
 * @param hostIP The IP address of the host
 * @param hostVulnMap The host vulnerability map containing severity counts
 * @returns The vulnerability severity counts for the host, or a zeroed object if not found
 */
export function getHostVulnerabilityCounts(
  hostIP: string,
  hostVulnMap: Record<string, HostVulnerabilityCounts>
): HostVulnerabilityCounts {
  // Return the host's vulnerability counts if available
  if (hostVulnMap && hostVulnMap[hostIP]) {
    return hostVulnMap[hostIP];
  }

  // If not found, return zeroed counts
  return {
    critical: 0,
    high: 0,
    medium: 0,
    low: 0,
    informational: 0,
    total: 0,
  };
}

/**
 * Creates a vulnerability map from the aggregated vulnerability response
 * @param aggregatedData The aggregated vulnerability data
 * @returns A map from host IPs to their vulnerability counts
 */
export function createHostVulnerabilityMap(
  aggregatedData: VulnerabilityAggregationResponse
): HostVulnerabilityMap {
  const map = new Map<string, HostVulnerabilityCounts>();

  // Process host totals
  if (aggregatedData && aggregatedData.hostTotals) {
    Object.entries(aggregatedData.hostTotals).forEach(([ip, counts]) => {
      map.set(ip, counts);
    });
  }

  return map;
}

/**
 * Calculates a risk score for a host based on vulnerability severity counts
 * @param counts The host's vulnerability severity counts
 * @returns A risk score from 0-100
 */
export function calculateHostRiskScore(
  counts: HostVulnerabilityCounts
): number {
  if (!counts || counts.total === 0) {
    return 0;
  }

  // Weight each severity level
  const weights = {
    critical: 10,
    high: 5,
    medium: 2,
    low: 1,
    informational: 0.2,
  };

  // Calculate weighted score
  const score =
    counts.critical * weights.critical +
    counts.high * weights.high +
    counts.medium * weights.medium +
    counts.low * weights.low +
    counts.informational * weights.informational;

  // Cap at 100 and round to nearest integer
  return Math.min(Math.round(score), 100);
}
